'''
Created on May 31, 2009

@author: mkiyer
'''

import math
import numpy

def write_wig_fixed_step(fhd, chrom, cov_array, span=None):
    # reduce entire chromosome to only the nonzero elements
    a = numpy.flatnonzero(cov_array)
    if len(a) == 0:
        print 'write_wig_fixed_step: no nonzero elements in array'
        return
        
    prev_i = a[0]
    fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, prev_i+1))
    fhd.write("%d\n" % cov_array[prev_i])    
    for i in a[1:]:
        if i != prev_i + 1:
            fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, i+1))            
        fhd.write("%d\n" % cov_array[i])
        prev_i = i
    return

def write_wig_fixed_step_v2(fhd, chrom, covarray, span=None):
    # reduce entire chromosome to only the nonzero elements
    a = numpy.flatnonzero(covarray)
    if len(a) == 0:
        print 'write_wig_fixed_step: no nonzero elements in array'
        return
    
    # cluster indices into groups of adjacent elements
    start = a[0]
    for i in xrange(1, len(a)-1):
        if a[i] + 1 != a[i+1]:
            end = a[i]
            
        
    print 'finding starts'
    
    starts = numpy.r_[0, numpy.array([a[i] for i in xrange(1, len(a)) if a[i]-1 != a[i-1]])]
    
    print 'finding ends'
    
    ends = numpy.r_[numpy.array([a[i] for i in xrange(0, len(a)-1) if a[i]+1 != a[i+1]]), a[-1]]
              
    
    print 'printing regions'
    #starts = numpy.r_[0, numpy.array([y for x,y in zip(a[:-1], a[1:]) if x+1 != y])]
    #ends = numpy.r_[numpy.array([x for x,y in zip(a[:-1], a[1:]) if x+1 != y]), a[-1]]
    assert len(starts) == len(ends)
    print len(starts)
    
    for chrom_start, chrom_end in zip(starts, ends):
        fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, chrom_start+1))
        covstr = map(str, covarray[chrom_start:chrom_end])
        covstr = '\n'.join(covstr)
        fhd.write(covstr)

    return

#    WRITING = 0
#    SEARCHING = 1
#    state = SEARCHING
#
#    for i,v in enumerate(cov_array):
#        if v != 0:
#            if state == SEARCHING: 
#                fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, i))    
#            fhd.write("%d\n" % v)
#            state = WRITING
#        else:
#            state = SEARCHING
#    start = 0
#    fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, start))    
#    for i,v in enumerate(cov_array):
#        fhd.write("%d\n" % v)    
#    return


def write_wig(fhd, name, chrom, cov_array):
    fhd.write("track type=wiggle_0 name=\"%s\"\n" % (name))
    # reset state
    cur_v = None
    span = None
    for i,v in enumerate(cov_array):
        if not cur_v:
            # reset state
            start, cur_v = i, v
            continue

        if v != cur_v:
            # if the interval value is zero, no need to write to the file
            # because the default value is zero
            if cur_v != 0:
                # save previous span and compute new one
                prev_span = span
                span = i - start

                if prev_span != span:                      
                    # mode change, span change, 
                    # or no mode was previously set
                    fhd.write("variableStep\tchrom=%s\tspan=%d\n" % (chrom, span))
                fhd.write("%d\t%f\n" % (start, cur_v * 10.0))
            # reset state
            start, cur_v = i, v

def write_wig_greedy(fhd, name, chrom, cov_array):
    '''
    implements greedy algorithm for choosing whether to write in variable 
    step or fixed step modes
    
    unfortunately the greedy algorithm always chooses variableStep
    '''
    fhd.write("track type=wiggle_0 name=\"%s\"\n" % (name))
    
    # approximate size of header lines for different modes 
    variableStep_header = 26
    fixedStep_header = 31

    # define modes
    FIXED_MODE = 0
    VARIABLE_MODE = 1
    
    cur_v = None
    mode = None
    span = None

    for i,v in enumerate(cov_array):
        if not cur_v:
            # reset state
            start, cur_v = i, v
            continue

        if v != cur_v:
            # if the interval value is zero, no need to write to the file
            # because the default value is zero
            if cur_v != 0:
                # save previous span and compute new one
                prev_span = span
                span = i - start
                # save previous mode and determine new mode to use
                prev_mode = mode
                # calculate cost of using fixedStep mode
                fixed_cost = (fixedStep_header + len(chrom) + len(str(start))) if mode == VARIABLE_MODE else 0
                fixed_cost += span * (len(str(cur_v)) + 1)
                # calculate cost of using variableStep mode
                variable_cost = (variableStep_header + len(chrom) + len(str(span))) if ((mode == FIXED_MODE) or (prev_span != span)) else 0 
                variable_cost += len(str(start)) + len(str(cur_v)) + 2
                if fixed_cost < variable_cost:
                    mode = FIXED_MODE
                else:
                    mode = VARIABLE_MODE

                if mode == VARIABLE_MODE:
                    if (prev_mode != mode) or (prev_span != span):                      
                        # mode change, span change, 
                        # or no mode was previously set
                        fhd.write("variableStep chrom=%s span=%d\n" % (chrom, span))
                    fhd.write("%d\t%d\n" % (start, cur_v))
                elif mode == FIXED_MODE:
                    if prev_mode != mode:
                        # mode change
                        fhd.write("fixedStep chrom=%s start=%d step=1\n" % (chrom, start))                        
                    for x in xrange(span):
                        fhd.write("%d\n" % cur_v)
            # reset state
            start, cur_v = i, v

class WIGFileWriter(object):
    '''
    Class to help create WIG files
    '''
    
    def __init__(self):
        pass
    
    def addTrack(self, **kwargs):
        '''
        Add the track definition line to the WIG file

        name              trackLabel           # default is "User Track"
        description       centerlabel          # default is "User Supplied Track"
        visibility        full|dense|hide      # default is hide (will also take numeric values 2|1|0)
        color             RRR,GGG,BBB          # default is 255,255,255
        altColor          RRR,GGG,BBB          # default is 128,128,128
        priority          N                    # default is 100
        autoScale         on|off               # default is off
        gridDefault       on|off               # default is off
        maxHeightPixels   max:default:min      # default is 128:128:11
        graphType         bar|points           # default is bar
        viewLimits        lower:upper          # default is range found in data
        yLineMark         real-value           # default is 0.0
        yLineOnOff        on|off               # default is off
        windowingFunction maximum|mean|minimum # default is maximum
        smoothingWindow   off|[2-16]           # default is off
        '''
        
        self.track_line = 'track type=wiggle_0 '
        for k, v in kwargs:
            self.track_line += k + '=' + str(v)
        pass



class WIGFile(object):
    '''
    Parser for wiggle (WIG) files
    
    https://cgwb.nci.nih.gov/goldenPath/help/wiggle.html
     
    track type=wiggle_0 name=track_label description=center_label \
        visibility=display_mode color=r,g,b altColor=r,g,b \
        priority=priority autoScale=on|off \
        gridDefault=on|off maxHeightPixels=max:default:min \
        graphType=bar|points viewLimits=lower:upper \
        yLineMark=real-value yLineOnOff=on|off \
        windowingFunction=maximum|mean|minimum smoothingWindow=off|2-16
    '''

    # parser states
    TRACK = 0
    VARIABLE_STEP = 1
    FIXED_STEP = 2
    BED = 3

    def __init__(self, fhd):
        '''
        Constructor
        '''
        self.fhd = fhd
        self.state = self.TRACK

    def parseTrack(self, stmt):
        pass
    
    def parseVariableStep(self):
        pass
    
    def parseFixedStep(self):
        pass
    
    @staticmethod
    def mapLog(thisline):
        value = float(thisline.strip())
        if value > 0:
            value = math.log(value, 20)
            #value = math.log10(value)
        elif value < 0:
            #value = -(math.log10(abs(value)))
            value = -(math.log(abs(value), 20))
        return str(value) + '\n'
    
    def map(self, func, outfhd):
        for thisline in self.fhd:
            if thisline.startswith('track'):
                # TODO: extract fields from track stmt
                pass
            elif thisline.startswith('fixedStep'):
                self.state = self.FIXED_STEP
                # TODO: extract fields from fixed step stmt
            elif thisline.startswith('variableStep'):
                self.state = self.VARIABLE_STEP
                # TODO: extract fields from variable step stmt
            else:
                if self.state == self.FIXED_STEP:
                    thisline = func(thisline)
                elif self.state == self.VARIABLE_STEP:
                    pass
                else:
                    # BED format
                    pass
            outfhd.write(thisline)